from typing import Dict, List, Tuple, Any


class GameState:
    """游戏状态管理器 - 管理单个游戏的状态信息"""
    
    def __init__(self):
        self.episode_reward = 0.0
        self.episode_length = 0
        self.episode_invalid_moves = 0
        self.transitions: List[Tuple] = []
    
    def update(self, state: Dict, action: int, reward: float, 
              next_state: Dict, done: bool, info: Dict) -> None:
        """更新游戏状态"""
        # 统计非法落子
        if info.get('invalid_move', False):
            self.episode_invalid_moves += 1
        
        # 存储transition
        transition = (
            self._state_to_array(state),
            action,
            reward,
            self._state_to_array(next_state),
            done
        )
        self.transitions.append(transition)
        
        # 更新统计
        self.episode_reward += reward
        self.episode_length += 1
    
    def get_results(self, winner: int) -> Tuple[float, int, int, int]:
        """获取游戏结果"""
        return self.episode_reward, self.episode_length, winner, self.episode_invalid_moves
    
    def _state_to_array(self, state: Dict) -> Any:
        """将环境状态转换为数组"""
        return state['board'].flatten()